import os
import numpy as np
import cv2
from concurrent.futures import ThreadPoolExecutor


def upsample_flow(flow, scale_factor, interpolation=cv2.INTER_CUBIC):
    h, w = flow.shape[:2]
    new_h, new_w = int(h * scale_factor), int(w * scale_factor)
    flow_up_x = cv2.resize(flow[:, :, 0], (new_w, new_h), interpolation=interpolation) * scale_factor
    flow_up_y = cv2.resize(flow[:, :, 1], (new_w, new_h), interpolation=interpolation) * scale_factor
    flow_up = np.stack((flow_up_x, flow_up_y), axis=2)
    return flow_up


def process_and_save_upsampled_flow(filename, input_dir, output_dir, scale_factor):
    input_flow_path = os.path.join(input_dir, filename)
    flow = np.load(input_flow_path)
    flow_upsampled = upsample_flow(flow, scale_factor)

    output_flow_path = os.path.join(output_dir, filename)
    np.save(output_flow_path, flow_upsampled)
    print(f"Upsampled and saved {filename} to {output_flow_path}")


def main(input_dir, output_dir, scale_factor, max_workers=40):
    os.makedirs(output_dir, exist_ok=True)
    flow_files = [f for f in os.listdir(input_dir) if f.endswith('.npy')]

    with ThreadPoolExecutor(max_workers=max_workers) as executor:
        futures = [executor.submit(process_and_save_upsampled_flow, filename, input_dir, output_dir, scale_factor)
                   for filename in flow_files]

    print(f"All flows upsampled and saved to {output_dir}")


# 示例用法
input_dir = '/root/data/deep_align/32nm_flow(3)'  # 替换为实际输入目录
output_dir = '/root/data/deep_align/32nm_flow_align_to_8nm'  # 替换为实际输出目录

# 上采样因子
scale_factor = 4096 / 1024

main(input_dir, output_dir, scale_factor)
